import math
import numbers

import pandas as pd
from sqlalchemy import create_engine, Engine
from pandas.core.frame import DataFrame
from sqlalchemy.orm import sessionmaker
from sqlalchemy.sql import text
from sqlalchemy import sql
from tools.tool import tool


class DbTool:
    _engine = None

    def getEngine(self) -> Engine:
        if self._engine is not None:
            return self._engine
        if tool.isTestEnv():
            engine = create_engine('mysql+pymysql://root:fw12345678@106.53.211.236:63306/stock')
        else:  # win=笔记本
            engine = create_engine('mysql+pymysql://root:fw12345678@localhost:3306/stock')
        self._engine = engine
        return engine

    def getDbConn(self):
        engine = self.getEngine()
        return engine.connect()

    def query(self, sql):
        with self.getDbConn() as conn:
            conn.execute(self.text(sql))
            conn.commit()

    def text(self, sql):
        if isinstance(sql, str):
            sql = text(sql)
        return sql

    def getDF(self, sql, indexCol=None):
        with dbTool.getDbConn() as conn:
            df = pd.read_sql(self.text(sql), conn, index_col=indexCol)
        # records = df.to_dict(orient='records')
        return df

    def getList(self, sql):
        df = self.getDF(sql)
        records = df.to_dict(orient='records')
        return records

    def getOne(self, sql):
        df = self.getDF(sql)
        return df.iloc[0]

    def save(self, data):
        Session = sessionmaker(bind=dbTool.getEngine())
        with Session() as session:
            session.merge(data)  # 无记录添加,有记录更新
            session.commit()
        # return obj

    def saveOne(self, table, data, tableComment=None):
        return self.saveAll(table, [data], True)

    def sqlAlchemySaveAll(self, table, data: list, getId=False):
        if not data:
            return 0
        # 假设你已经定义好连接
        connection = self.getDbConn()
        # 将data中的nan 替换为 None
        for d in data:
            for k, v in d.items():
                if isinstance(v, float) and math.isnan(v):
                    d[k] = None
        cols = ''
        colsPlaceholder = ''
        for key in data[0].keys():
            cols += '`' + key + '`' + ','
            colsPlaceholder += ':' + key + ','
        cols = cols[:-1]
        colsPlaceholder = colsPlaceholder[:-1]
        # 构造 SQL 语句和执行
        insert_statement = text(f"""
            REPLACE INTO {table} ({cols}) VALUES ({colsPlaceholder})
        """)
        res = connection.execute(insert_statement, data)
        connection.commit()
        if getId:
            return res.lastrowid
        else:
            return res.rowcount

    def saveAll(self, table, data, tableComment=None):
        """
        保存数据,如果表不存在创建数据库表
        """
        cnt = 0
        if isinstance(data, DataFrame):
            dictList = data.to_dict(orient='records')
            dataFrame = data
        else:
            dictList = data
            dataFrame = pd.DataFrame(data)
        try:
            cnt = self.sqlAlchemySaveAll(table, dictList)
        except BaseException as e1:
            if "doesn't exist" not in str(e1):
                raise e1
            with self.getEngine().connect() as conn:
                dataFrame.to_sql(table, con=conn, index=False, if_exists='append')
                print(f'可能数据表不存在,尝试使用pandas建表,并插入数据, 建表:{table}, 注意添加索引')
                # 修改表注释
                if tableComment is not None:
                    sql = f"alter table {table} comment '{tableComment}'"
                    conn.execute(text(sql))
                # exit(0)

        return cnt


dbTool = DbTool()
